import argparse
import os
import sys
from typing import Optional, Tuple

if sys.version_info < (3, 8):
    from typing_extensions import Literal
else:
    from typing import Literal

from github import Github
from openai import OpenAI
from tenacity import retry, stop_after_attempt, wait_random_exponential

GITHUB_TOKEN = os.environ.get("GITHUB_API_TOKEN")
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")

client = OpenAI(api_key=OPENAI_API_KEY)

Model = Literal["gpt-4", "gpt-3.5-turbo", "vicuna-7b-v1.1"]

CC_TYPES = os.environ.get(
    "CC_TYPES",
    ", ".join(
        [
            "feat",
            "fix",
            "docs",
            "style",
            "refactor",
            "perf",
            "test",
            "build",
            "ci",
            "chore",
            "revert",
            "security",
        ]
    ),
)

CC_SCOPES = os.environ.get(
    "CC_SCOPES",
    ", ".join(
        [
            "sdk",
            "cli",
            "public-api",
            "integrations",
            "artifacts",
            "media",
            "sweeps",
            "launch",
        ]
    ),
)


@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
def chat_completion_with_backoff(**kwargs):
    """Call OpenAI's chat completion API with exponential backoff."""
    return client.chat.completions.create(**kwargs)


@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
def get_pr_info(
    pr_number: int,
    repo_name: str = "wandb/wandb",
    get_diff: bool = True,
) -> Tuple[str, Optional[str]]:
    """Get the title and diff of a PR."""
    g = Github(GITHUB_TOKEN)
    repo = g.get_repo(repo_name)

    # Get specific pull request by number
    pr = repo.get_pull(pr_number)

    if not get_diff:
        return pr.title, None

    # Get the files
    files = pr.get_files()

    diff = "\n".join(
        [
            file.patch or file.filename  # use filename if patch is empty
            for file in files
            if "_pb2" not in file.filename  # ignore autogenerated protobuf files
        ]
    )
    pr_title = pr.title

    return pr_title, diff


def check_pr_title(
    pr_number: int,
    model: Model = "gpt-4",
    repo_name: str = "wandb/wandb",
) -> bool:
    """Check whether a PR title follows the conventional commit format."""
    pr_title, _ = get_pr_info(pr_number, repo_name=repo_name, get_diff=False)

    messages = [
        {
            "role": "system",
            "content": (
                "Your task is to check whether the title for a GitHub pull request "
                "follows the conventional commit format: "
                f"<type>(<scope>): <description>. The possible types are: {CC_TYPES}."
                f"The possible scopes are: {CC_SCOPES}."
                "The description must start with a verb in the imperative mood and be lower case."
                "The user will provide the current PR title. You must respond with 'yes' or 'no'."
            ),
        },
        {
            "role": "user",
            "content": f"{pr_title}",
        },
    ]

    response = chat_completion_with_backoff(
        model=model,
        messages=messages,
    )
    # fixme:
    is_compliant = response.choices[0]["message"]["content"].lower().strip()
    print(is_compliant)

    return is_compliant == "yes"


def generate_pr_title(
    pr_number: int,
    model: Model = "gpt-4",
    repo_name: str = "wandb/wandb",
) -> str:
    """Generate a PR title for a given PR number using the given model."""
    messages = [
        {
            "role": "system",
            "content": (
                "Your task is to write a title for a GitHub pull request "
                "that will follow the conventional commit format and capture the essence of the change: "
                f"<type>(<scope>): <description>. The possible types are: {CC_TYPES}."
                f"The possible scopes are: {CC_SCOPES}."
                "The description must start with a verb in the imperative mood and be lower case."
                "For context, the user will provide the current title and the diff of the pull request."
                "and their corresponding labels."
                "You must respond in the format: <type>(<scope>): <description>."
                "Be concise and specific. If you are unsure, keep the original title."
                "Even if you think the correct type or scope is missing, you must only use the provided options."
                "Be concise."
            ),
        },
        {
            "role": "user",
            "content": "Title: {{TITLE}}. Diff: {{DIFF}}.",
        },
    ]

    pr_title, diff = get_pr_info(pr_number, repo_name=repo_name)

    # todo: check context limit, strip diff if too long

    messages[-1]["content"] = messages[-1]["content"].replace("{{TITLE}}", pr_title)
    messages[-1]["content"] = messages[-1]["content"].replace("{{DIFF}}", diff)

    completion = chat_completion_with_backoff(
        model=model,
        messages=messages,
    )
    suggested_title = completion.choices[0]["message"]["content"]

    return suggested_title


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # add two subparsers: one for the "generate" command and one for the "check" command
    subparsers = parser.add_subparsers(dest="command")

    generate_parser = subparsers.add_parser("generate", help="Generate PR title")
    generate_parser.add_argument("pr_number", type=int, help="Pull Request number")

    check_parser = subparsers.add_parser("check", help="Check PR title")
    check_parser.add_argument("pr_number", type=int, help="Pull Request number")

    args = parser.parse_args()

    if args.command == "generate":
        title = generate_pr_title(args.pr_number)
        print(title)
    elif args.command == "check":
        is_conventional_commit_compliant = check_pr_title(args.pr_number)
        if not is_conventional_commit_compliant:
            raise ValueError(
                "PR title is not compliant with the conventional commit recommendations. \n"
                "Comment on your PR with `/suggest-title` to get a suggestion or "
                "`/fix-title` to ask the pr-title-bot to fix it for you."
            )
    else:
        print("Invalid command. Use 'generate' or 'check'")
